package com.easydata.agent.service;

import com.easydata.agent.consts.TaskState;
import com.easydata.agent.core.DAGContext;
import com.easydata.agent.core.Task;
import com.easydata.agent.core.TaskPool;
import com.easydata.agent.core.entity.TaskDefinition;
import com.easydata.agent.exceptions.TaskException;
import com.easydata.agent.util.Base64Utils;
import com.easydata.libPlugin.core.AbstractDAGNode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.Base64;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Queue;

@Slf4j
@Service
public class TaskService {
    @Autowired
    DAGService dagService;
    @Autowired
    TaskPool taskPool;

    public Task createTask(TaskDefinition definition) throws TaskException {

        var task = new Task();
        task.setId(definition.getId());
        task.setState(TaskState.CREATED);
        task.setTaskDefinition(definition);
        task.setContext(new DAGContext());
        task.getContext().setTaskId(task.getId());
        task.setDag(dagService.buildDAG(task.getContext(), definition));
        task.getContext().setParameters(definition.getParameter());

        return task;
    }

    public void prepareTask(Task task) throws Exception {

        if (task.getState() != TaskState.CREATED) {
            throw new TaskException(-1, "task is not in CREATED state");
        }

        for (AbstractDAGNode n : task.getDag().getNodes()) {
            n.prepare(task.getContext());
        }

        task.setState(TaskState.PREPARED);
        taskPool.addTask(task);
    }

    public void executeTask(Task task) throws Exception {
        if (task.getState() != TaskState.PREPARED) {
            throw new TaskException(-1, "task is not in PREPARED state");
        }

        List<AbstractDAGNode> list = dagService.findEntry(task.getDag());
        // fill the incoming data to first level nodes
        if (!task.getTaskDefinition().getData().isEmpty()) {
            for (var node : list) {
                node.getWriter().open();
                node.getWriter().write(Base64Utils.decode(task.getTaskDefinition().getData()));
                node.getWriter().close();
            }
        }

        Queue<AbstractDAGNode> visitQueue = new PriorityQueue<>();
        visitQueue.addAll(list);

        // BFS
        // level by level visit
        while (!visitQueue.isEmpty()) {
            AbstractDAGNode node = visitQueue.poll();
            node.getReader().open();
            node.getWriter().open();
            node.start(task.getContext());
            node.process(task.getContext());
            node.stop(task.getContext());
            node.getReader().close();
            node.getWriter().close();

            var backwardNodes = task.getDag().getBackwardNodes(node.getId());
            visitQueue.addAll(backwardNodes);
        }

        taskPool.removeTask(task.getId());
    }
}
